import json
import numpy as np
from utils.node import Node

def load_data(file_path = "data/testData3.json"):

    with open(file_path,"r") as f:
        json_data = json.load(f)

    nodes = json_data["nodes"]
    edges = json_data["edges"]
    load_data = {data_["rdfId"]:data_["load"] for data_ in json_data["loadDataList"]}
    equip_data = json_data["equipData"]


    generator_types = set([
        "Substation",
    ])
    switch_types = set([
        "GroundDisconnector","Disconnector","Breaker","LoadBreakSwitch", 
        "O/C","station_end",
    ])

    load_types = set(
        ["TransformerWinding","PowerTransformer"]
        )

    generators = {}
    loads = {}
    lines = {}

    to_maps = {}
    from_maps = {}
    for edge in edges:
        if edge["from"] not in to_maps:
            to_maps[edge["from"]] = [edge["to"],]
        else:
            to_maps[edge["from"]].append(edge["to"])
        if edge["to"] not in from_maps:
            from_maps[edge["to"]] = [edge["from"],]
        else:
            from_maps[edge["to"]].append(edge["from"])

    for node in nodes:
        node = Node(
            node["baseVoltage"],
            node["capacity"],
            node["open"],
            node["rdfId"],
            node["type"],
            node["virtualNode"],
            from_maps[node["rdfId"]] if node["rdfId"] in from_maps else [],
            to_maps[node["rdfId"]] if node["rdfId"] in to_maps else [],
        )
        if node["node_type"] in generator_types:
            generators[node["rdf_id"]] = node
        elif node["node_type"] in load_types:
            node["load"] = load_data[node["rdf_id"]]
            loads[node["rdf_id"]] = node

        else:
            if node["node_type"] in switch_types:
                node["capacity"] = None
                node["rated_capacity"] = None
                
            lines[node["rdf_id"]] = node
    del nodes
    return generators, loads, lines




def buid_test_data(generators_config = 
                   [
                       ["G1",500],
                       ["G2",500],
                   ],
                   loads_config = [
                       ["U1",80],
                       ["U2",130],
                       ["U3",90],
                   ],
                   points_config = [
                       ["A", None],
                       ["B", None],
                       ["C", None],
                       ["D", None],
                       ["E", None],
                       ["F", None],
                       ["G", None],
                   ],
                   lines_config = [
                       [["G1","A"],300],
                       [["G1","B"],200],
                       [["G2","G"],100],
                       [["A","B"],200],
                       [["A","C"],200],
                       [["B","C"],200],
                    #    [["B","U3"],200],
                       [["C","D"],90],
                       [["C","F"],200],
                       [["G","F"],200],
                       [["G","U3"],100],
                       [["F","D"],200],
                       [["D","U1"],200],
                       [["F","U2"],200],
                   ],
                   ):
        
        nodes = {}

        generators = {}
        for config_ in generators_config:
            node_ = Node(
                base_voltage="",
                capacity=config_[1],
                is_open=False,
                rdf_id=config_[0],
                node_type="Substation",
                is_virtual_node=False,
                from_list=[],
                to_list=[]
            )
            generators[config_[0]] = node_
            nodes[config_[0]] = node_

        loads = {}
        for config_ in loads_config:
            node_ = Node(
                base_voltage="",
                capacity=np.inf,
                is_open=False,
                rdf_id=config_[0],
                node_type="TransformerWinding",
                is_virtual_node=False,
                from_list=[],
                to_list=[]
            )
            node_["load"] = [config_[1]]
            loads[config_[0]] = node_
            nodes[config_[0]] = node_

        lines={}
        for config_ in points_config:
            node_ = Node(
                base_voltage="",
                capacity=config_[1],
                is_open=False,
                rdf_id=config_[0],
                node_type="Fuse",
                is_virtual_node=False,
                from_list=[],
                to_list=[]
            )
            lines[config_[0]] = node_
            nodes[config_[0]] = node_
            
        for line_config in lines_config:
            point_from = line_config[0][0]
            point_to = line_config[0][1]
            cap = line_config[1]
            rdf_id= "{}_{}".format(point_from, point_to)

            node_ = Node(
                base_voltage="",
                capacity=cap,
                is_open=False,
                rdf_id=rdf_id,
                node_type="ACLineSegment",
                is_virtual_node=False,
                from_list=[point_from],
                to_list=[point_to])
            lines[rdf_id] = node_
            nodes[rdf_id] = node_

            nodes[point_from]["to_list"].append(rdf_id)
            nodes[point_to]["from_list"].append(rdf_id)

        return generators,loads,lines
            
            
                
            
        


        #         ('G1', 'A'): 80, 
        # ('G1', 'B'): 200, 
        # ('A', 'U1'): 200, 
        # ('B', 'A'): 200, 


if __name__ == "__main__":
    # generators, loads, lines = buid_test_data()
    generators, loads, lines = load_data()
    import pickle
    with open("nodes_data.pickle", "rb") as f:
        nodes = pickle.load(f)
    for load_id, load in loads.items() :
        load: Node
        line:list = nodes[load_id]["best_path"][::-1]
        assert line.pop(0) == load_id # 确认负载是对的
        assert line.pop(-1) in generators # 确认发电设备是对的
        assert line[0] in load.get_from_list() # 确认线路正常连接到负载
        # 遍历整个线路是否完整
        for pre_node_idx_, node_idx_ in enumerate(line[1:]):
            assert line[pre_node_idx_] in lines[node_idx_]
            assert node_id_ in lines[line[idx_-1]].get_from_list()
            
    print("done")